#include "compi/topology.h"

#include <stdexcept>

namespace compi
{
namespace topology
{

TopologyHandle::TopologyHandle(Context & ctx, MPI_Comm comm, bool owns) noexcept
    : _ctx(&ctx), _comm(comm), _owns(owns)
{
}

TopologyHandle::~TopologyHandle()
{
    reset();
}

TopologyHandle::TopologyHandle(TopologyHandle && other) noexcept
{
    _ctx = other._ctx;
    _comm = other._comm;
    _owns = other._owns;
    other._ctx = nullptr;
    other._comm = MPI_COMM_NULL;
    other._owns = false;
}

TopologyHandle &
TopologyHandle::operator=(TopologyHandle && other) noexcept
{
    if (this != &other)
    {
        reset();
        _ctx = other._ctx;
        _comm = other._comm;
        _owns = other._owns;
        other._ctx = nullptr;
        other._comm = MPI_COMM_NULL;
        other._owns = false;
    }
    return *this;
}

MPI_Comm
TopologyHandle::release() noexcept
{
    MPI_Comm tmp = _comm;
    _comm = MPI_COMM_NULL;
    _owns = false;
    _ctx = nullptr;
    return tmp;
}

void
TopologyHandle::reset() noexcept
{
    if (_owns && _ctx && _comm != MPI_COMM_NULL)
    {
        (void)_ctx->call_mpi([&]() -> int { return MPI_Comm_free(&_comm); },
                             "MPI_Comm_free (TopologyHandle)");
    }
    _comm = MPI_COMM_NULL;
    _owns = false;
    _ctx = nullptr;
}

namespace
{
bool
validate_cart_params(const CartCreateParams & params) noexcept
{
    return !params.dims.empty() && params.dims.size() == params.periods.size();
}
bool
validate_graph_params(const GraphCreateParams & params, int comm_size) noexcept
{
    if (params.index.empty())
        return false;
    if (static_cast<int>(params.index.size()) != comm_size)
        return false;
    if (params.index.back() != static_cast<int>(params.edges.size()))
        return false;
    return true;
}
} // namespace

Result<TopologyHandle>
create_cartesian(Context & ctx, const CartCreateParams & params) noexcept
{
    if (!validate_cart_params(params))
    {
        return Result<TopologyHandle>::error(
            ResultCode::InvalidRequest,
            "Invalid CartCreateParams: dims/periods must be non-empty and have the same size.");
    }

    MPI_Comm base = (params.base_comm == MPI_COMM_NULL) ? ctx.comm() : params.base_comm;
    MPI_Comm new_comm = MPI_COMM_NULL;

    auto cart_res = ctx.call_mpi(
        [&]() -> int
        { return MPI_Cart_create(base,
                                 static_cast<int>(params.dims.size()),
                                 params.dims.data(),
                                 params.periods.data(),
                                 params.reorder ? 1 : 0,
                                 &new_comm); },
        "MPI_Cart_create");

    if (cart_res.is_error())
    {
        return Result<TopologyHandle>::error(cart_res.error_code(), cart_res.error_message());
    }

    TopologyHandle handle(ctx, new_comm, new_comm != MPI_COMM_NULL);
    return Result<TopologyHandle>::success(std::move(handle));
}

Result<std::vector<int>>
cart_coords(const TopologyHandle & topo, int rank) noexcept
{
    if (!topo.valid() || topo.context() == nullptr)
    {
        return Result<std::vector<int>>::error(ResultCode::InvalidRequest,
                                               "TopologyHandle is not valid");
    }

    Context & ctx = *topo.context();
    int ndims = 0;
    auto ndims_res = ctx.call_mpi(
        [&]() -> int { return MPI_Cartdim_get(topo.comm(), &ndims); }, "MPI_Cartdim_get");
    if (ndims_res.is_error())
    {
        return Result<std::vector<int>>::error(ndims_res.error_code(), ndims_res.error_message());
    }

    std::vector<int> coords(static_cast<std::size_t>(ndims), 0);
    auto coords_res = ctx.call_mpi(
        [&]() -> int { return MPI_Cart_coords(topo.comm(), rank, ndims, coords.data()); },
        "MPI_Cart_coords");
    if (coords_res.is_error())
    {
        return Result<std::vector<int>>::error(coords_res.error_code(), coords_res.error_message());
    }

    return Result<std::vector<int>>::success(std::move(coords));
}

Result<int>
cart_rank(const TopologyHandle & topo, const std::vector<int> & coords) noexcept
{
    if (!topo.valid() || topo.context() == nullptr)
    {
        return Result<int>::error(ResultCode::InvalidRequest, "TopologyHandle is not valid");
    }
    Context & ctx = *topo.context();

    int ndims = 0;
    auto ndims_res = ctx.call_mpi(
        [&]() -> int { return MPI_Cartdim_get(topo.comm(), &ndims); }, "MPI_Cartdim_get");
    if (ndims_res.is_error())
    {
        return Result<int>::error(ndims_res.error_code(), ndims_res.error_message());
    }

    if (static_cast<int>(coords.size()) != ndims)
    {
        return Result<int>::error(ResultCode::InvalidRequest,
                                  "Coordinate dimension mismatch with cart communicator");
    }

    int rank = 0;
    auto rank_res = ctx.call_mpi(
        [&]() -> int { return MPI_Cart_rank(topo.comm(), coords.data(), &rank); }, "MPI_Cart_rank");
    if (rank_res.is_error())
    {
        return Result<int>::error(rank_res.error_code(), rank_res.error_message());
    }

    return Result<int>::success(rank);
}

Result<TopologyHandle>
create_graph(Context & ctx, const GraphCreateParams & params) noexcept
{
    MPI_Comm base = (params.base_comm == MPI_COMM_NULL) ? ctx.comm() : params.base_comm;
    int size = 0;
    auto size_res = ctx.call_mpi([&]() -> int { return MPI_Comm_size(base, &size); }, "MPI_Comm_size");
    if (size_res.is_error())
    {
        return Result<TopologyHandle>::error(size_res.error_code(), size_res.error_message());
    }

    if (!validate_graph_params(params, size))
    {
        return Result<TopologyHandle>::error(ResultCode::InvalidRequest,
                                             "Invalid GraphCreateParams: index/edges mismatch or size mismatch.");
    }

    MPI_Comm graph_comm = MPI_COMM_NULL;
    auto graph_res = ctx.call_mpi(
        [&]() -> int
        {
            return MPI_Graph_create(base,
                                    size,
                                    params.index.data(),
                                    params.edges.data(),
                                    params.reorder ? 1 : 0,
                                    &graph_comm);
        },
        "MPI_Graph_create");

    if (graph_res.is_error())
    {
        return Result<TopologyHandle>::error(graph_res.error_code(), graph_res.error_message());
    }

    TopologyHandle handle(ctx, graph_comm, graph_comm != MPI_COMM_NULL);
    return Result<TopologyHandle>::success(std::move(handle));
}

Result<std::vector<int>>
graph_neighbors(const TopologyHandle & topo, int max_neighbors) noexcept
{
    if (!topo.valid() || topo.context() == nullptr)
    {
        return Result<std::vector<int>>::error(ResultCode::InvalidRequest,
                                               "TopologyHandle is not valid");
    }
    Context & ctx = *topo.context();

    int nneighbors = 0;
    auto neigh_res = ctx.call_mpi(
        [&]() -> int { return MPI_Graph_neighbors_count(topo.comm(), ctx.rank(), &nneighbors); },
        "MPI_Graph_neighbors_count");
    if (neigh_res.is_error())
    {
        return Result<std::vector<int>>::error(neigh_res.error_code(), neigh_res.error_message());
    }

    if (max_neighbors > 0 && nneighbors > max_neighbors)
    {
        nneighbors = max_neighbors;
    }

    std::vector<int> neighbors(static_cast<std::size_t>(nneighbors));
    auto neighbors_res = ctx.call_mpi(
        [&]() -> int { return MPI_Graph_neighbors(topo.comm(), ctx.rank(), nneighbors, neighbors.data()); },
        "MPI_Graph_neighbors");
    if (neighbors_res.is_error())
    {
        return Result<std::vector<int>>::error(neighbors_res.error_code(), neighbors_res.error_message());
    }

    return Result<std::vector<int>>::success(std::move(neighbors));
}

} // namespace topology
} // namespace compi
